package org.momo.rpc.registry.zookeeper;

import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.momo.rpc.config.config.registry.ZkRegistryConfig;
import org.momo.rpc.constant.Constant;
import org.momo.rpc.registry.ConsumerRegistry;
import org.momo.rpc.registry.ProviderRegistry;
import org.momo.rpc.provider.model.ProviderRegistryInfo;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;


@Slf4j
public class ZkRegistry implements ConsumerRegistry, ProviderRegistry {

    private static final Set<String> registeredService = ConcurrentHashMap.newKeySet();

    private static final Set<String> registeredRetryMethod = ConcurrentHashMap.newKeySet();

    private static final Map<String, List<ProviderRegistryInfo>> cachedServiceList = new ConcurrentHashMap<>();

    @Override
    public void initRegistryClient(ZkRegistryConfig config) throws Exception {
        Curator.initClient(config);
    }

    @Override
    public List<ProviderRegistryInfo> discover(String serviceName, String serviceVersion, String serviceKey) {
        // select service from local cache
        List<ProviderRegistryInfo> possibleServiceList = cachedServiceList.get(serviceKey);

        // no cache, fetch from zookeeper
        if (possibleServiceList == null || possibleServiceList.isEmpty()) {
            List<String> childrenNodePath = Curator.getChildrenNode(serviceName, serviceVersion);
            if (childrenNodePath != null && !childrenNodePath.isEmpty()) {
                possibleServiceList = convert(serviceName, serviceVersion, childrenNodePath);
                cachedServiceList.put(serviceKey, possibleServiceList);
                addListenerForService(serviceName, serviceVersion);
            }
        }
        return possibleServiceList;
    }

    @Override
    public boolean canRetry(String serviceName, String serviceVersion, String methodName) {
        String serviceKey = String.join(Constant.FUNCTION_DELIMITER, serviceName, serviceVersion, methodName);
        if(registeredRetryMethod.contains(serviceKey)) {
            return true;
        }
        if (Curator.checkNodeExist(serviceName, serviceVersion, methodName)) {
            registeredRetryMethod.add(serviceKey);
            return true;
        }
        return false;
    }

    @Override
    public void register(ProviderRegistryInfo providerRegistryInfo) {
        // service has been registered
        String serviceKey = providerRegistryInfo.getServiceKey();
        if (registeredService.contains(serviceKey)) {
            return;
        }

        // register to zookeeper
        Curator.createNode(providerRegistryInfo.getServiceName(), providerRegistryInfo.getServiceVersion(), providerRegistryInfo.getServiceAddress());
        registeredService.add(serviceKey);
    }

    @Override
    public void registerRetryMethod(String serviceName, String serviceVersion, List<String> retryMethods) {
        for(String retryMethod : retryMethods) {
            Curator.createNode(Constant.RETRY, serviceName, serviceVersion, retryMethod);
            String sign = String.join(Constant.FUNCTION_DELIMITER, serviceName, serviceVersion, retryMethod);
            registeredRetryMethod.add(sign);
        }
    }

    @Override
    public void unRegister(ProviderRegistryInfo providerRegistryInfo) {
        String serviceId = providerRegistryInfo.getServiceKey();
        if (!registeredService.contains(serviceId)) {
            log.error("Service {} has not been registered", serviceId);
            return;
        }
        Curator.removeNode(providerRegistryInfo.getServiceName(), providerRegistryInfo.getServiceVersion(), providerRegistryInfo.getHost(), providerRegistryInfo.getPort().toString());
        registeredService.remove(serviceId);
    }

    @Override
    public void closeRegistryClient() {
        Curator.closeClient();
    }

    private List<ProviderRegistryInfo> convert(String serviceName, String serviceVersion, List<String> childrenNodePath) {
        return childrenNodePath.stream()
                .map((child) -> ProviderRegistryInfo.convertToMetaInfo(serviceName, serviceVersion, child))
                .collect(Collectors.toList());
    }

    private void addListenerForService(String serviceName, String serviceVersion) {
        String finalServiceName = String.join(Constant.ZOOKEEPER_DELIMITER, "", serviceName, serviceVersion);
        PathChildrenCacheListener listener = (client, event) -> {
            List<String> serviceList = client.getChildren().forPath(finalServiceName);
            cachedServiceList.put(finalServiceName, convert(serviceName, serviceVersion, serviceList));
        };
        Curator.addListener(finalServiceName, listener);
    }

    private void addListenerForRetry(String serviceName, String serviceVersion) {
        String finalServiceName = String.join(Constant.ZOOKEEPER_DELIMITER, "", serviceName, serviceVersion);
        PathChildrenCacheListener listener = (client, event) -> {
            List<String> serviceList = client.getChildren().forPath(finalServiceName);
            cachedServiceList.put(finalServiceName, convert(serviceName, serviceVersion, serviceList));
        };
        Curator.addListener(finalServiceName, listener);
    }

}
